package com.example.h9iserver.service.springAi.rag;

import org.springframework.ai.autoconfigure.vectorstore.pgvector.PgVectorStoreProperties;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;
import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import javax.sql.DataSource;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

@Configuration
@EnableConfigurationProperties(PgVectorStoreProperties.class)
public class PgVectorConfig {

    // 手动创建VectorStore Bean，指定使用postgres数据源
    @Bean
    @ConditionalOnMissingBean
    public VectorStore vectorStore(EmbeddingModel embeddingModel, DataSource dataSource) {
        // 从动态数据源中获取postgres数据源
        DynamicRoutingDataSource dynamicDataSource = (DynamicRoutingDataSource) dataSource;
        DataSource postgresDataSource = dynamicDataSource.getDataSource("postgres");

        // 使用postgres数据源创建JdbcTemplate
        JdbcTemplate jdbcTemplate = new JdbcTemplate(postgresDataSource);

        // 创建并返回PgVectorStore实例
        VectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, embeddingModel)
                .dimensions(1024)                    // Optional: defaults to model dimensions or 1536
                .distanceType(COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(true)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("vector_store")     // Optional: defaults to "vector_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
        return vectorStore;
    }


}
